# 导包
import pandas as pd
import datetime
import matplotlib.pyplot as plt
from utils.log import Logger
from utils.common import data_preprocessing
from xgboost import XGBRFClassifier
from sklearn.model_selection import train_test_split
# 导入随机森林库
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, roc_auc_score
from sklearn.metrics import accuracy_score
import joblib

# 解决图表中文乱码问题
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15


# 定义类，配置日志，获取数据源
class PowerLoadModel(object):
    def __init__(self, filename):
        # 配置日志记录
        logfile_name = "train_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        # 获取数据源
        self.data_source = data_preprocessing(filename)


# 数据分析处理
def ana_data(data):
    data = data.copy(deep=True)
    fig = plt.figure(figsize=(20, 32))
    # 1.离职与年龄的比例关系直方图
    plt.subplot(5, 1, 1)
    plt.hist(data[data['Attrition'] == 1]['Age'], bins=20, alpha=0.5, label='离职')
    plt.legend()
    plt.title('离职与年龄比例关系直方图')
    plt.xlabel('年龄')
    plt.ylabel('人数')
    # 自动调整子图间距
    plt.tight_layout()
    #  2.离职与学历的比例关系直方图
    plt.subplot(5, 1, 2)
    plt.hist(data[data['Attrition'] == 1]['Education'], bins=20, alpha=0.5, label='离职')
    plt.legend()
    # 将X轴的值映射为对应的学历名称 1.0大学以下  2.0学士  3.0硕士  4.0博士
    plt.xticks([1.0, 2.0, 3.0, 4.0, 5.0], ['大学以下', '大学', '学士', '硕士', '博士'])
    plt.title('离职与学历比例关系直方图')
    plt.xlabel('学历')
    plt.ylabel('人数')
    # 3.离职与职位的比例关系扇形比例图
    plt.subplot(5, 1, 3)
    plt.pie(data[data['Attrition'] == 1]['JobRole'].value_counts(),
            labels=data[data['Attrition'] == 1]['JobRole'].value_counts().index, autopct='%.2f%%')
    plt.title('离职与职位比例关系扇形比例图')
    plt.xlabel('职位')
    plt.ylabel('人数')
    # 4.月工资与离职的比例关系折线图  人数为y轴,月工资为x轴
    # 提前筛选离职员工数据，提升可读性和效率
    left_employees = data[data['Attrition'] == 1]
    # 按照月工资进行分箱，每500为一个区间（可根据实际数据调整）
    income_bins = pd.cut(left_employees['MonthlyIncome'], bins=range(0, int(data['MonthlyIncome'].max()) + 500, 500))
    # 统计每个工资区间的离职人数
    income_grouped = left_employees.groupby(income_bins).size()
    # 绘图
    plt.subplot(5, 1, 4)
    plt.plot(income_grouped.index.map(lambda x: x.mid), income_grouped.values, marker='o')
    plt.title('月工资与离职比例关系折线图')
    plt.xlabel('月工资（区间中点）')
    plt.ylabel('离职人数')
    #  5.离职与工作年限折现图
    grouped = left_employees.groupby('YearsAtCompany')['Attrition'].sum().reset_index()
    plt.subplot(5, 1, 5)
    plt.plot(grouped['YearsAtCompany'], grouped['Attrition'], marker='o')
    plt.title('离职与工作年限折现图')
    plt.xlabel('工作年限')
    plt.ylabel('离职人数')
    plt.tight_layout()  # 自动调整子图间距
    plt.savefig('../data/fig/负荷分析图.png')


# 特征分析
def feature_ana(data, logger):
    logger.info("===============开始进行特征工程处理===============")
    result = data.copy(deep=True)
    logger.info("===============开始进行特征工程处理===============")
    # 1.获取特征数据 x删除Attrition，EmployeeNumber，Over18，StandardHours，PerformanceRating，EnvironmentSatisfaction
    x = data.drop(columns=['Attrition', 'EmployeeNumber', 'Over18', 'StandardHours', 'PerformanceRating',
                           'EnvironmentSatisfaction', 'Gender', 'Department', 'StockOptionLevel'])
    # 假设 df 是你的原始 DataFrame，包含特征和标签 Department
    categorical_cols = ['BusinessTravel', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']
    # # 对分类变量进行 One-Hot 编码
    x = pd.get_dummies(x, columns=categorical_cols)
    #  分离特征和标签
    print(x.columns)
    x = x.drop(columns=['OverTime_No'], axis=1)
    y = data['Attrition']
    return x, y


#  模型训练
def train_model(x, y, logger):
    logger.info("===============开始进行模型训练===============")
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=1)
    # 2.网格化搜索与交叉验证
    # 2.1备选的超参数
    # print("开始网格化搜索")
    # param_grid = {
    #     'max_depth': [10,20,30,40,50,60,100],
    #     # 'learning_rate': [0.1,0.5,0.7, 0.01,0.07],
    #     'n_estimators': [1,5,10,20,50,70, 100],
    #     'min_samples_split': [2, 3, 4, 5, 6, 7, 8, 9, 10],
    # }

    # # 1. 设置参数搜索空间（示例）
    # param_grid = {
    #     'C': [0.1, 1, 10],          # 正则化强度的倒数
    #     'penalty': ['l1', 'l2'],     # 正则化类型
    #     'solver': ['liblinear']      # 支持L1/L2的求解器
    # }
    #
    # # 2. 使用LogisticRegression作为估计器
    # estimator = LogisticRegression(max_iter=1000)
    #
    # # 2.2 实例化网格搜索，配置交叉验证
    # grid_search = GridSearchCV(
    #     estimator=estimator,
    #     param_grid=param_grid,
    #     cv=5,
    #     scoring='accuracy',         # 明确指定评分方式
    #     n_jobs=-1,                  # 启用多线程加速
    #     verbose=1                   # 输出详细过程
    # )
    #
    # grid_search.fit(x_train, y_train)
    # print("最佳参数：", grid_search.best_params_)
    # print("最佳模型：", grid_search.best_estimator_)
    # y_pre = grid_search.predict(x_test)
    # print("准确率：", accuracy_score(y_test, y_pre))
    # print(f"AUC：{roc_auc_score(y_test, y_pre)}")
    # xgb = XGBRFClassifier(max_depth=20, learning_rate=0.5, n_estimators=42)
    # xgb = RandomForestClassifier(max_depth=40, min_samples_split=3, n_estimators=10)
    xgb = LogisticRegression(C=10, max_iter=1000, penalty='l1', solver='liblinear')
    xgb.fit(x_train, y_train)
    # 4.模型评价
    # 4.1模型在训练集上的预测结果
    y_pred_train = xgb.predict(x_train)
    # 4.2模型在测试集上的预测结果
    y_pred_test = xgb.predict(x_test)
    # 4.3模型在训练集上的MSE、MAPE
    mse_train = mean_squared_error(y_true=y_train, y_pred=y_pred_train)
    mae_train = mean_absolute_error(y_true=y_train, y_pred=y_pred_train)
    print(f"模型在训练集上的均方误差：{mse_train}")
    print(f"模型在训练集上的平均绝对误差：{mae_train}")
    # 4.4模型在测试集上的MSE、MAPE
    mse_test = mean_squared_error(y_true=y_test, y_pred=y_pred_test)
    mae_test = mean_absolute_error(y_true=y_test, y_pred=y_pred_test)
    print(f"模型在测试集上的均方误差：{mse_test}")
    print(f"模型在测试集上的平均绝对误差：{mae_test}")
    print(f"AUC：{roc_auc_score(y_test, y_pred_test)}")
    logger.info("=========================模型训练完成=============================")
    logger.info(f"模型在训练集上的均方误差：{mse_train}")
    logger.info(f"模型在训练集上的平均绝对误差：{mae_train}")
    logger.info(f"模型在测试集上的均方误差：{mse_test}")
    logger.info(f"模型在测试集上的平均绝对误差：{mae_test}")
    joblib.dump(xgb, '../model/xgb.pkl')

if __name__ == '__main__':
    df = PowerLoadModel('../data/train.csv')
    ana_data(df.data_source)
    x, y = feature_ana(df.data_source, df.logfile)
    train_model(x, y, df.logfile)
